# Copyright (c) Facebook, Inc. and its affiliates.
import json
import os

from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.utils.file_io import PathManager
from detectron2.data.datasets import load_sem_seg

UAVID_8_CATEGORIES = [
    {"color": [0, 0, 0], "id": 0, "isthing": 0, "name": "Building"},
    {"color": [1, 1, 1], "id": 1, "isthing": 0, "name": "Road"},
    {"color": [2, 2, 2], "id": 2, "isthing": 0, "name": "Tree"},
    {"color": [3, 3, 3], "id": 3, "isthing": 0, "name": "LowVeg"},
    {"color": [4, 4, 4], "id": 4, "isthing": 0, "name": "Moving_Car"},
    {"color": [5, 5, 5], "id": 5, "isthing": 0, "name": "Static_Car"},
    {"color": [6, 6, 6], "id": 6, "isthing": 0, "name": "Human"},
    {"color": [7, 7, 7], "id": 7, "isthing": 0, "name": "Clutter"},
    # {"color": [255, 255, 255], "id": 7, "isthing": 0, "name": "Boundary"},
]
UAVID_SEM_SEG_CATEGORIES = [
    "Building",
    "Road",
    "Tree",
    "LowVeg",
    "Moving_Car",
    "Static_Car",
    "Human",
    "Clutter"
]
UAVID_COLORS = [k["color"] for k in UAVID_8_CATEGORIES]

MetadataCatalog.get("uavid_sem_seg_train").set(
    stuff_colors=UAVID_COLORS[:],
)

MetadataCatalog.get("uavid_sem_seg_val").set(
    stuff_colors=UAVID_COLORS[:],
)

def register_sem_seg_uavid(root):
    root = os.path.join(root, "uavid")
    root = "/root/autodl-tmp/data/uavid"
    for name, dirname in [("train", "train_512"), ("val", "val_512")]:
        image_dir = os.path.join(root, dirname, "images")
        gt_dir = os.path.join(root, dirname, "masks")
        name = f"uavid_sem_seg_{name}"
        DatasetCatalog.register(
            name, lambda x=image_dir, y=gt_dir: load_sem_seg(y, x, gt_ext="png", image_ext="png")
        )
        MetadataCatalog.get(name).set(
            stuff_classes=UAVID_SEM_SEG_CATEGORIES[:],
            image_root=image_dir,
            sem_seg_root=gt_dir,
            evaluator_type="sem_seg",
            ignore_label=255,
        )
_root = os.getenv("DETECTRON2_DATASETS", "datasets")
register_sem_seg_uavid(_root)
